import pandas as pd
import numpy as np
import lightgbm as lgb
from sklearn.model_selection import train_test_split
from sklearn.metrics import mean_absolute_error
import matplotlib.pyplot as plt

# 加载数据
data = pd.read_csv("../demo/historical_sales_data.csv")

# 将日期转换为 datetime 类型
data['date_id'] = pd.to_datetime(data['date_id'])


# 按 SKU 分组处理
def preprocess_data(df):
    # 添加时间特征
    df['dayofweek'] = df['date_id'].dt.dayofweek  # 星期几
    df['dayofmonth'] = df['date_id'].dt.day  # 每月第几天
    df['month'] = df['date_id'].dt.month  # 月份

    # 添加滞后特征
    for lag in [1, 2, 3, 7, 14]:
        df[f'sales_lag_{lag}'] = df['sales'].shift(lag)

    # 添加滑动平均特征
    df['sales_rolling_7'] = df['sales'].rolling(window=7).mean()

    # 删除缺失值
    df = df.dropna()
    return df


# 对每个 SKU 进行预处理
data = data.groupby('sku_id').apply(preprocess_data,include_groups=False).reset_index(drop=True)


# 特征列
features = ['dayofweek', 'dayofmonth', 'month', 'price', 'stock', 'is_promotion',
            'sales_lag_1', 'sales_lag_2', 'sales_lag_3', 'sales_lag_7', 'sales_lag_14',
            'sales_rolling_7']

# 目标变量
target = 'sales'

# 划分训练集和测试集
train_data = data[data['date_id'] < '2023-08-15']
test_data = data[data['date_id'] >= '2023-08-15']

X_train = train_data[features]
y_train = train_data[target]
X_test = test_data[features]
y_test = test_data[target]





# 创建 LightGBM 数据集
train_set = lgb.Dataset(X_train, label=y_train)

# 定义模型参数
params = {
    'objective': 'regression',
    'metric': 'mae',
    'boosting_type': 'gbdt',
    'num_leaves': 31,
    'learning_rate': 0.05,
    'feature_fraction': 0.9,
    'verbose': -1
}

# 训练模型
model = lgb.train(params, train_set, num_boost_round=100)



################################
# 预测测试集
y_pred = model.predict(X_test)

# 将预测结果添加到测试集
test_data['pred_sales'] = y_pred

# 输出未来7天的预测结果
future_predictions = test_data[['date_id', 'sku_id', 'sales', 'pred_sales']]
print(future_predictions.tail(7))



################################
# 绘制实际销量与预测销量对比图
plt.figure(figsize=(12, 6))
plt.plot(test_data['date_id'], test_data['sales'], label='Actual Sales')
plt.plot(test_data['date_id'], test_data['pred_sales'], label='Predicted Sales', linestyle='--')
plt.title('Actual vs Predicted Sales')
plt.xlabel('date_id')
plt.ylabel('Sales')
plt.legend()
plt.show()